import os

import casadi
import meshcat.geometry as mg
import numpy as np
import pinocchio as pin
from gymnasium import spaces
from pinocchio import casadi as cpin
from pinocchio.visualize import MeshcatVisualizer
from scipy.spatial.transform import Rotation as R

# serve: pin.RobotWrapper.BuildFromURDF
os.environ["ROS_PACKAGE_PATH"] = "piper_description"


class Arm_IK:
    """
    reduced_robot:
    robot
    pin.forwardKinematics(model, data) will modify the data in-replace

    FYI:
    len(data.joints) = len(data.oMi)
    len(data.oMf) gives info of every body
    """

    def __init__(self):
        np.set_printoptions(precision=5, suppress=True, linewidth=200)
        urdf_dir = "piper_description/piper_description.urdf"
        self.robot = pin.RobotWrapper.BuildFromURDF(urdf_dir)
        self.mixed_jointsToLockIDs = ["joint7", "joint8"]
        self.reduced_robot = self.robot.buildReducedRobot(
            list_of_joints_to_lock=self.mixed_jointsToLockIDs,
            reference_configuration=np.array([0] * self.robot.model.nq),
        )

        q = R.from_euler("xyz", (0, 0, 0)).as_quat()  # (x, y, z, w)
        self.reduced_robot.model.addFrame(
            pin.Frame(
                "effector",
                self.reduced_robot.model.getJointId("joint6"),  #
                pin.SE3(
                    pin.Quaternion(q),  # (x, y, z, w)
                    np.array([0.0, 0.0, 0.14]),  # shift on z-axis
                ),
                pin.FrameType.OP_FRAME,
            )
        )

        self.geom_model = pin.buildGeomFromUrdf(self.robot.model, urdf_dir, pin.GeometryType.COLLISION)
        for i in range(4, 9):
            for j in range(0, 3):
                self.geom_model.addCollisionPair(pin.CollisionPair(i, j))
        self.geometry_data = pin.GeometryData(self.geom_model)

        self.init_data = np.zeros(self.reduced_robot.model.nq)
        self.history_data = np.zeros(self.reduced_robot.model.nq)
        # Creating Casadi models and data for symbolic computing
        self.cmodel = cpin.Model(self.reduced_robot.model)
        self.cdata = self.cmodel.createData()

        # Creating symbolic variables and define the loss function
        self.cq = casadi.SX.sym("q", self.reduced_robot.model.nq, 1)  # the joints
        self.cTf = casadi.SX.sym("tf", 4, 4)  # the transformation matrix for the target position
        cpin.framesForwardKinematics(self.cmodel, self.cdata, self.cq)
        self.gripper_id = self.reduced_robot.model.getFrameId("effector")
        self.loss = casadi.Function(
            "loss",
            [self.cq, self.cTf],
            [
                casadi.vertcat(
                    cpin.log6(self.cdata.oMf[self.gripper_id].inverse() * cpin.SE3(self.cTf)).vector,
                )
            ],
        )
        """
        If only cares about the position residual:
            cpin.log6(self.cdata.oMf[self.gripper_id].inverse() * cpin.SE3(self.cTf)).vector[:3]
        If only cares about the orientation residual:
            cpin.log6(self.cdata.oMf[self.gripper_id].inverse() * cpin.SE3(self.cTf)).vector[3:]
        
        Above: self.cdata.oMf[self.gripper_id] gives the abstract SE3 of the gripper is of size [4, 4]:
        example of a valued se3:
        array([ [-0.82956, -0.23614,  0.50603,  0.35228],
                [-0.14509,  0.96621,  0.21303,  0.2687 ],
                [-0.53924,  0.1033 , -0.83579,  0.0111 ],
                [ 0.     ,  0.     ,  0.     ,  1.     ]])
                
        self.cTf corresponds to the target.homogeneous in the world frame:
        array([[-0.82687, -0.23449,  0.51118,  0.34803],
            [-0.14329,  0.96678,  0.21169,  0.27306],
            [-0.54383,  0.10179, -0.833  ,  0.     ],
            [ 0.     ,  0.     ,  0.     ,  1.     ]])
            
        Here the multiplication inside log6 is:
        
        pin.SE3(end_effector.homogeneous).inverse() * pin.SE3(target.homogeneous)
        
        This gives a transformation: Origion -> target -> "Try go origion using effector".
        
        How to measure "go back to origion"?
        
        cpin.log6(SE3) gives the Lie algebra of the transformation, which is accidentally a 6D-vector for each SE3.
        
        Here the objective is to minimize the L2-norm of the Lie aglebra(6D-vector) so that the pin.SE3(end_effector.homogeneous) 
        is almost the inverse of pin.SE3(target.homogeneous) by adjusting the var_q (joints). 
        If perfectly solved:
        
        SE3(end_effector.homogeneous) = SE3(target.homogeneous).inverse()
        
        Fetch the value?
        use self.opti.value(?)
        you will find that:
        (almost) self.opti.value(self.total_loss) == sum((pin.log((ee_SE3).inverse() * tar_SE3).vector)**2)
        
        """

        # Defining the optimization problem
        self.opti = casadi.Opti()
        self.var_q = self.opti.variable(self.reduced_robot.model.nq)
        self.param_tf = self.opti.parameter(4, 4)
        self.total_loss = casadi.sumsqr(self.loss(self.var_q, self.param_tf))
        self.regularization = casadi.sumsqr(self.var_q)

        # Setting optimization constraints and goals
        self.opti.subject_to(
            self.opti.bounded(
                self.reduced_robot.model.lowerPositionLimit, self.var_q, self.reduced_robot.model.upperPositionLimit
            )
        )
        self.opti.minimize(20 * self.total_loss + 0.001 * self.regularization)

        opts = {"ipopt": {"print_level": 0, "max_iter": 1000, "tol": 1e-6}, "print_time": False}
        self.opti.solver("ipopt", opts)

        self.joint_space = spaces.Box(
            low=self.reduced_robot.model.lowerPositionLimit.astype(np.float32),
            high=self.reduced_robot.model.upperPositionLimit.astype(np.float32),
            shape=(self.reduced_robot.model.nq,),
        )

    def forwardKinematics(self, q) -> pin.SE3:
        data = self.reduced_robot.model.createData()
        pin.forwardKinematics(self.reduced_robot.model, data, q)
        pin.updateFramePlacements(self.reduced_robot.model, data)
        return data.oMf[self.gripper_id]

    def visualize(self, sol_q: np.ndarray, target_xyz: np.ndarray, target_quat: np.ndarray):
        """
        sol_q: the solution joint in radian
        target_pose: target (x, y, z, rx, ry, rz)
        """

        # q = R.from_euler("xyz", (rx, ry, rz)).as_quat()  # (x, y, z, w)
        target = pin.SE3(
            pin.Quaternion(target_quat),  # (x, y, z, w)
            target_xyz,
        )

        # Initialize the Meshcat visualizer  for visualization
        self.vis = MeshcatVisualizer(
            self.reduced_robot.model, self.reduced_robot.collision_model, self.reduced_robot.visual_model
        )
        self.vis.initViewer(open=True)
        self.vis.loadViewerModel("pinocchio")
        self.vis.displayFrames(True, frame_ids=[113, 114], axis_length=0.15, axis_width=5)
        self.vis.display(pin.neutral(self.reduced_robot.model))

        # to visualize the position of ee
        self.vis.viewer["ee_sphere"].set_object(
            mg.Sphere(0.003), mg.MeshPhongMaterial(color=0x0000FF, transparent=True, opacity=0.7)
        )  # blue
        self.vis.viewer["tar_sphere"].set_object(
            mg.Sphere(0.003), mg.MeshPhongMaterial(color=0xFF0000, transparent=True, opacity=0.7)
        )  # red

        # a slice of cylinder to indicate the direction of the target (gripper catch!)
        arrow_geometry = mg.Cylinder(height=0.001, radius=0.01)
        arrow_material = mg.MeshPhongMaterial(color=0x00FF00, transparent=True, opacity=0.5)

        # 设置箭头对象
        self.vis.viewer["target_arrow"].set_object(arrow_geometry, arrow_material)

        # Enable the display of end effector target frames with short axis lengths and greater width.
        frame_viz_names = ["ee_frame"]
        FRAME_AXIS_POSITIONS = (
            np.array([[0, 0, 0], [1, 0, 0], [0, 0, 0], [0, 1, 0], [0, 0, 0], [0, 0, 1]]).astype(np.float32).T
        )
        FRAME_AXIS_COLORS = (
            np.array([[1, 0, 0], [1, 0.6, 0], [0, 1, 0], [0.6, 1, 0], [0, 0, 1], [0, 0.6, 1]]).astype(np.float32).T
        )
        for frame_viz_name in frame_viz_names:
            self.vis.viewer[frame_viz_name].set_object(
                mg.LineSegments(
                    mg.PointsGeometry(
                        position=0.1 * FRAME_AXIS_POSITIONS,
                        color=FRAME_AXIS_COLORS,
                    ),
                    mg.LineBasicMaterial(
                        linewidth=20,
                        vertexColors=True,
                    ),
                )
            )

        self.vis.display(sol_q)
        # for visualization of the frame axis: only requires transforms
        # set transform: use the homogenous matrix
        # show the target frame
        self.vis.viewer["tar_sphere"].set_transform(target.homogeneous)
        self.vis.viewer["target_arrow"].set_transform(target.homogeneous)
        # visualize the ee: requires forward kinematics
        data = self.reduced_robot.model.createData()
        pin.forwardKinematics(self.reduced_robot.model, data, sol_q)
        pin.updateFramePlacements(self.reduced_robot.model, data)
        ee_pose = data.oMf[self.gripper_id]

        # print("ee SE3", ee_pose)
        self.vis.viewer["ee_sphere"].set_transform(ee_pose.homogeneous)
        self.vis.viewer["ee_frame"].set_transform(ee_pose.homogeneous)

    def ik_solve(self, target_xyz: np.ndarray, target_quat: np.ndarray, gripper=0):

        # q = R.from_euler("xyz", (roll, pitch, yaw)).as_quat()
        target = pin.SE3(
            pin.Quaternion(target_quat),  # (x, y, z, w)
            target_xyz,
        )
        # print("target SE3:",target)

        gripper = np.array([gripper / 2.0, -gripper / 2.0])
        # if motorstate is not None:
        #     self.init_data = motorstate
        # self.opti.set_initial(self.var_q, self.init_data)
        self.opti.set_initial(self.var_q, self.init_data)
        self.opti.set_value(self.param_tf, target.homogeneous)
        # self.opti.set_value(self.var_q_last, self.init_data) # for smooth

        sol = self.opti.solve()
        # sol = self.opti.solve_limited()
        sol_q = self.opti.value(self.var_q)

        is_collision = self.check_self_collision(sol_q, gripper)

        # get other data
        data = self.reduced_robot.model.createData()
        pin.forwardKinematics(self.reduced_robot.model, data, sol_q)
        pin.updateFramePlacements(self.reduced_robot.model, data)
        sol_pose = data.oMf[self.gripper_id]

        sol_pos = sol_pose.translation
        sol_euler = pin.rpy.matrixToRpy(sol_pose.rotation)
        sol_x = np.concatenate([sol_pos, sol_euler], axis=-1)

        info = {
            "loss": self.opti.value(self.total_loss),
            "regularization": self.opti.value(self.regularization),
            "sol_q": sol_q,
            "sol_x": sol_x,
            "collision": is_collision,
            "sol_state": sol,
        }

        return sol_q, info

    def check_self_collision(self, q, gripper=np.array([0, 0])):
        """
        check whether the arm collide with itself
        """
        pin.forwardKinematics(self.robot.model, self.robot.data, np.concatenate([q, gripper], axis=0))
        pin.updateGeometryPlacements(self.robot.model, self.robot.data, self.geom_model, self.geometry_data)
        collision = pin.computeCollisions(self.geom_model, self.geometry_data, False)
        # print("collision:", collision)
        return collision


if __name__ == "__main__":
    """
    target pos in meters!
    target rot in rad !
    """
    self = Arm_IK()

    # t_pos = (0.34803, 0.27306, 0, 3.12, 0.575, -2.97)
    # q = R.from_euler("xyz", (3.14, 0, -3.14)).as_quat()
    # t_pos = (0.34803, 0.27306, 0, q)
    # t_pos = [0.19561, -0.0, 0.22539, -0.00004, 1.48359, -0.00005]

    xyz = np.array([0.34803, 0.27306, 0])
    quat = np.array([0.08522, -0.95511, -0.03461, -0.28161])
    # quat = np.array([0, -1, 0, 0])

    sol_q, info = self.ik_solve(xyz, quat)

    self.visualize(sol_q, xyz, quat)

    rand_q1 = self.joint_space.sample()
    rand_q2 = self.joint_space.sample()

    data = self.reduced_robot.model.createData()
    pin.forwardKinematics(self.reduced_robot.model, data, rand_q1)
    pin.updateFramePlacements(self.reduced_robot.model, data)

    start_se3 = data.oMf[self.gripper_id]
    start_pos = start_se3.translation
    start_euler = pin.rpy.matrixToRpy(start_se3.rotation)

    start_x = np.concatenate([start_pos, start_euler], axis=-1)

    data = self.reduced_robot.model.createData()
    pin.forwardKinematics(self.reduced_robot.model, data, rand_q2)
    pin.updateFramePlacements(self.reduced_robot.model, data)

    target_se3 = data.oMf[self.gripper_id]
    target_pos = start_se3.translation
    target_euler = pin.rpy.matrixToRpy(start_se3.rotation)

    target_x = np.concatenate([start_pos, start_euler], axis=-1)
